{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84053a41-1280-4c3f-9122-c418cd6740e1",
   "metadata": {},
   "source": [
    "# Import packages and read variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6c8351c6-6dd5-4180-a8e5-cb8fd4e75f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "import google.generativeai as genai\n",
    "import google.ai.generativelanguage as glm\n",
    "from elasticsearch import Elasticsearch, helpers\n",
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "GOOGLE_API_KEY = os.getenv(\"GOOGLE_API_KEY\")\n",
    "ES_USER = os.getenv(\"ES_USER\")\n",
    "ES_PASSWORD = os.getenv(\"ES_PASSWORD\")\n",
    "elastic_index_name='gemini-demo'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad658ee0-f9cc-4e15-91cd-c32a2e00ec44",
   "metadata": {},
   "source": [
    "# Connnet to Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a05950ff-6593-43e9-875f-d4338f8a9f4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'name': 'liuxgm.local', 'cluster_name': 'elasticsearch', 'cluster_uuid': 'p4U3azj3TbSz7Fm_9W3Ylg', 'version': {'number': '8.12.0', 'build_flavor': 'default', 'build_type': 'tar', 'build_hash': '1665f706fd9354802c02146c1e6b5c0fbcddfbc9', 'build_date': '2024-01-11T10:05:27.953830042Z', 'build_snapshot': False, 'lucene_version': '9.9.1', 'minimum_wire_compatibility_version': '7.17.0', 'minimum_index_compatibility_version': '7.0.0'}, 'tagline': 'You Know, for Search'}\n"
     ]
    }
   ],
   "source": [
    "url = f\"https://{ES_USER}:{ES_PASSWORD}@192.168.0.3:9200\"\n",
    "\n",
    "es = Elasticsearch(\n",
    "        hosts=[url], \n",
    "        ca_certs = \"./http_ca.crt\", \n",
    "        verify_certs = True\n",
    ")\n",
    "print(es.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d9a776c-f1f4-497a-ad3b-32a92676f06b",
   "metadata": {},
   "source": [
    "# Delete index if it exists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aacf9c8f-5a70-4728-a589-a7276b693ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "if(es.indices.exists(index=elastic_index_name)):\n",
    "    print(\"The index has already existed, going to remove it\")\n",
    "    es.options(ignore_status=404).indices.delete(index=elastic_index_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad6b0b1a-0dc5-42c3-952f-879a2681c93c",
   "metadata": {},
   "source": [
    "# Index document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3a34724d-30bb-4ac8-9b58-be1ca93a32fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'_index': 'gemini-demo', '_id': 'hee2JY0BdnWNZCoqwm9Q', '_version': 1, 'result': 'created', '_shards': {'total': 2, 'successful': 1, 'failed': 0}, '_seq_no': 0, '_primary_term': 1}\n"
     ]
    }
   ],
   "source": [
    "genai.configure(api_key=GOOGLE_API_KEY)\n",
    "\n",
    "title = \"Beijing\"\n",
    "sample_text = (\"Beijing is the capital of China and the center of Chinese politics, culture, and economy. This city has a long history with many ancient buildings and cultural heritage. Beijing is renowned as a cultural city in China, boasting numerous museums, art galleries, and historical landmarks. Additionally, as a modern metropolis, Beijing is a thriving business center with modern architecture and advanced transportation systems. It serves as the seat of the Chinese government, where significant decisions and events often take place. Overall, Beijing holds a crucial position in China, serving as both a preserver of traditional culture and a representative of modern development.\")\n",
    "\n",
    "model = 'models/embedding-001'\n",
    "embedding = genai.embed_content(model=model,\n",
    "                                content=sample_text,\n",
    "                                task_type=\"retrieval_document\",\n",
    "                                title=title)\n",
    "\n",
    "doc = {\n",
    "    'text' : sample_text,\n",
    "    'text_embedding' : embedding['embedding'] \n",
    "}\n",
    "\n",
    "resp = es.index(index=elastic_index_name, document=doc)\n",
    "\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "713e935f-d15c-4271-af9f-f222038b06dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'_index': 'gemini-demo', '_id': 'hue3JY0BdnWNZCoqfG9w', '_version': 1, 'result': 'created', '_shards': {'total': 2, 'successful': 1, 'failed': 0}, '_seq_no': 1, '_primary_term': 1}\n"
     ]
    }
   ],
   "source": [
    "title = \"Shanghai\"\n",
    "sample_text = (\"Shanghai is one of China's largest cities and a significant hub for economy, finance, and trade. This modern city is located in the eastern part of China and serves as an international metropolis. The bustling streets, skyscrapers, and modern architecture in Shanghai showcase the city's prosperity and development. As one of China's economic engines, Shanghai is home to the headquarters of many international companies and various financial institutions. It is also a crucial trading port, connecting with destinations worldwide. Additionally, Shanghai boasts a rich cultural scene, including art galleries, theaters, and historical landmarks. In summary, Shanghai is a vibrant, modern city with international influence.\")\n",
    "\n",
    "model = 'models/embedding-001'\n",
    "embedding = genai.embed_content(model=model,\n",
    "                                content=sample_text,\n",
    "                                task_type=\"retrieval_document\",\n",
    "                                title=title)\n",
    "\n",
    "doc = {\n",
    "    'text' : sample_text,\n",
    "    'text_embedding' : embedding['embedding'] \n",
    "}\n",
    "\n",
    "resp = es.index(index=elastic_index_name, document=doc)\n",
    "\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed421f7d-31aa-4c19-8dce-1d68cdb43e4f",
   "metadata": {},
   "source": [
    "# Search for document with Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b10c5127-54cb-4b98-af16-0a18c699c75a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(question):\n",
    "    print(\"\\n\\nQuestion: \", question)\n",
    "    embedding = genai.embed_content(model=model,\n",
    "                                    content=question,\n",
    "                                    task_type=\"retrieval_query\")\n",
    "\n",
    "    resp = es.search(\n",
    "    index = elastic_index_name,\n",
    "    knn={\n",
    "        \"field\": \"text_embedding\",\n",
    "        \"query_vector\":  embedding['embedding'],\n",
    "        \"k\": 10,\n",
    "        \"num_candidates\": 100\n",
    "        }\n",
    "    )\n",
    "\n",
    "    for result in resp['hits']['hits']:\n",
    "        pretty_output = (f\"\\n\\nID: {result['_id']}\\n\\nText: {result['_source']['text']}\")\n",
    "        print(pretty_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "77c12a9b-4f52-433a-972f-6d24620fecbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Question:  How do you describe Beijing?\n",
      "\n",
      "\n",
      "ID: hee2JY0BdnWNZCoqwm9Q\n",
      "\n",
      "Text: Beijing is the capital of China and the center of Chinese politics, culture, and economy. This city has a long history with many ancient buildings and cultural heritage. Beijing is renowned as a cultural city in China, boasting numerous museums, art galleries, and historical landmarks. Additionally, as a modern metropolis, Beijing is a thriving business center with modern architecture and advanced transportation systems. It serves as the seat of the Chinese government, where significant decisions and events often take place. Overall, Beijing holds a crucial position in China, serving as both a preserver of traditional culture and a representative of modern development.\n",
      "\n",
      "\n",
      "ID: hue3JY0BdnWNZCoqfG9w\n",
      "\n",
      "Text: Shanghai is one of China's largest cities and a significant hub for economy, finance, and trade. This modern city is located in the eastern part of China and serves as an international metropolis. The bustling streets, skyscrapers, and modern architecture in Shanghai showcase the city's prosperity and development. As one of China's economic engines, Shanghai is home to the headquarters of many international companies and various financial institutions. It is also a crucial trading port, connecting with destinations worldwide. Additionally, Shanghai boasts a rich cultural scene, including art galleries, theaters, and historical landmarks. In summary, Shanghai is a vibrant, modern city with international influence.\n"
     ]
    }
   ],
   "source": [
    "search(\"How do you describe Beijing?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "342218ff-7fbb-45b4-8749-f624a1b7eb96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Question:  What is Shanghai like?\n",
      "\n",
      "\n",
      "ID: hue3JY0BdnWNZCoqfG9w\n",
      "\n",
      "Text: Shanghai is one of China's largest cities and a significant hub for economy, finance, and trade. This modern city is located in the eastern part of China and serves as an international metropolis. The bustling streets, skyscrapers, and modern architecture in Shanghai showcase the city's prosperity and development. As one of China's economic engines, Shanghai is home to the headquarters of many international companies and various financial institutions. It is also a crucial trading port, connecting with destinations worldwide. Additionally, Shanghai boasts a rich cultural scene, including art galleries, theaters, and historical landmarks. In summary, Shanghai is a vibrant, modern city with international influence.\n",
      "\n",
      "\n",
      "ID: hee2JY0BdnWNZCoqwm9Q\n",
      "\n",
      "Text: Beijing is the capital of China and the center of Chinese politics, culture, and economy. This city has a long history with many ancient buildings and cultural heritage. Beijing is renowned as a cultural city in China, boasting numerous museums, art galleries, and historical landmarks. Additionally, as a modern metropolis, Beijing is a thriving business center with modern architecture and advanced transportation systems. It serves as the seat of the Chinese government, where significant decisions and events often take place. Overall, Beijing holds a crucial position in China, serving as both a preserver of traditional culture and a representative of modern development.\n"
     ]
    }
   ],
   "source": [
    "search(\"What is Shanghai like?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "047d641e-b7da-4e6c-b9e3-11cff6f7180b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Question:  which city is the capital of China?\n",
      "\n",
      "\n",
      "ID: hee2JY0BdnWNZCoqwm9Q\n",
      "\n",
      "Text: Beijing is the capital of China and the center of Chinese politics, culture, and economy. This city has a long history with many ancient buildings and cultural heritage. Beijing is renowned as a cultural city in China, boasting numerous museums, art galleries, and historical landmarks. Additionally, as a modern metropolis, Beijing is a thriving business center with modern architecture and advanced transportation systems. It serves as the seat of the Chinese government, where significant decisions and events often take place. Overall, Beijing holds a crucial position in China, serving as both a preserver of traditional culture and a representative of modern development.\n",
      "\n",
      "\n",
      "ID: hue3JY0BdnWNZCoqfG9w\n",
      "\n",
      "Text: Shanghai is one of China's largest cities and a significant hub for economy, finance, and trade. This modern city is located in the eastern part of China and serves as an international metropolis. The bustling streets, skyscrapers, and modern architecture in Shanghai showcase the city's prosperity and development. As one of China's economic engines, Shanghai is home to the headquarters of many international companies and various financial institutions. It is also a crucial trading port, connecting with destinations worldwide. Additionally, Shanghai boasts a rich cultural scene, including art galleries, theaters, and historical landmarks. In summary, Shanghai is a vibrant, modern city with international influence.\n"
     ]
    }
   ],
   "source": [
    "search(\"which city is the capital of China?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "049ec37e-f8cd-46dc-b789-48bf220a3888",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Question:  the economy engine in China\n",
      "\n",
      "\n",
      "ID: hue3JY0BdnWNZCoqfG9w\n",
      "\n",
      "Text: Shanghai is one of China's largest cities and a significant hub for economy, finance, and trade. This modern city is located in the eastern part of China and serves as an international metropolis. The bustling streets, skyscrapers, and modern architecture in Shanghai showcase the city's prosperity and development. As one of China's economic engines, Shanghai is home to the headquarters of many international companies and various financial institutions. It is also a crucial trading port, connecting with destinations worldwide. Additionally, Shanghai boasts a rich cultural scene, including art galleries, theaters, and historical landmarks. In summary, Shanghai is a vibrant, modern city with international influence.\n",
      "\n",
      "\n",
      "ID: hee2JY0BdnWNZCoqwm9Q\n",
      "\n",
      "Text: Beijing is the capital of China and the center of Chinese politics, culture, and economy. This city has a long history with many ancient buildings and cultural heritage. Beijing is renowned as a cultural city in China, boasting numerous museums, art galleries, and historical landmarks. Additionally, as a modern metropolis, Beijing is a thriving business center with modern architecture and advanced transportation systems. It serves as the seat of the Chinese government, where significant decisions and events often take place. Overall, Beijing holds a crucial position in China, serving as both a preserver of traditional culture and a representative of modern development.\n"
     ]
    }
   ],
   "source": [
    "search(\"the economy engine in China\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80b9d51e-04d8-4a66-86df-c790aaa5b643",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python_3.11",
   "language": "python",
   "name": "python_3.11"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
